# -*- coding: utf-8 -*-

# -------------------------------
# Author    : linlingwei
# Email     : 951035650@qq.com
# File Name : deformerCore.py
# Date      : 2018/8/21
# IDE       : PyCharm
# Version   : 1.0.2
# -------------------------------


__author__ = "lin lingwei"

import os
import sys
import json
import time

import maya.cmds as mc
import maya.mel as mel
import pymel.core as pm

import scripts.commonCmds as cc


def findRelatedDeformer(deformerObj, typ="skinCluster"):
    """
    v1.0.1 : find related deformer.
    
    :param deformerObj: input maya object.
    :param typ: deformer type.
    :return: deformer name.
    """
    
    skinShapeWithPath = ""
    cp = mc.ls(deformerObj, type="controlPoint")
    if len(cp):
        deformerShape = deformerObj
    else:
        rels = mc.listRelatives(deformerObj)
        if not  rels:
            return
        for r in rels:
            if len(mc.ls(r, type="controlPoint")) == 0:
                continue
            if mc.getAttr(r + ".io"):
                continue
            if not mc.getAttr(r + ".v"):
                hiddenShape = r
                hiddenShapeWithPath = deformerObj + "|" + r
                continue
            deformerShape = r
            skinShapeWithPath = deformerObj + "|" + r
            break
    if len(deformerShape) == 0:
        if len(hiddenShape) == 0:
            return
        else:
            deformerShape = hiddenShape
            skinShapeWithPath = hiddenShapeWithPath
    if typ == "skinCluster":
        for cluster in mc.ls(type="skinCluster"):
            _geo = mc.skinCluster(cluster, q=True, g=True)
            if _geo:
                for geo in _geo:
                    if geo == deformerShape or geo == skinShapeWithPath:
                        return cluster
    else:
        clusterList = list()
        for cluster in mc.ls(type=typ):
            deformerGeo = mc.deformer(cluster, q=True, g=True)
            if deformerGeo:
                for geo in mc.deformer(cluster, q=True, g=True):
                    if geo == deformerShape or geo == skinShapeWithPath:
                        clusterList.append(cluster)
        return clusterList
    return

def skinCluterUnusedInf(skinNode):
    
    geo = mc.skinCluster(skinNode, q=True, g=True)
    pruneSmallWeight(geo)
    winfs = mc.skinCluster(skinNode, q=True, wi=True)
    infs = mc.skinCluster(skinNode, q=True, inf=True)
    unused_inf = list(set(infs) - set(winfs))
    
    return unused_inf

def removeUnusedSkinClusterInf(skinNode):
    """
    v1.2 : The code is cleaner.
    
    :param skinNode:
    :return:
    """
    unused_inf = skinCluterUnusedInf(skinNode)
    for inf in unused_inf:
        mc.skinCluster(skinNode, e=True, ri=inf)
    return

def exportSkinClusterWeight(skinObj, path, fileName):
    skinNode = findRelatedDeformer(skinObj)
    if not skinNode:
        return
    skinObj = pm.PyNode(skinNode)
    geo = mc.skinCluster(skinNode, q=True, g=True)[0]
    weight = skinObj.getWeights(geo)
    infJoint = mc.skinCluster(skinNode, q=True, inf=True)
    weight_dict = dict()
    weight_dict['joint'] = infJoint
    weight_value = dict()
    weight_dict['weight'] = weight_value
    
    for jot in infJoint:
        weight_value[jot] = []
    
    for vtx, wet in enumerate(weight):
        for i, w in enumerate(wet):
            weight_value[infJoint[i]].append(w)
    
    filePath = os.path.join(path, fileName)
    star = time.clock()
    
    with open(filePath, 'w') as fj:
        json.dump(weight_dict, fj, sort_keys=True, indent=4, separators=(',', ': '))
    
    end = time.clock()
    pm.displayInfo("Weight : Export skinCluster weight file path - {}".format(path))
    pm.displayInfo("Weight : Export skinCluster weight file name - {}".format(fileName + ".json"))
    pm.displayInfo("Weight : Export skinCluster elapsed time - {} s".format(end - star))
    
    return


def importSkinWeight(skinObj, path, fileName):
    skinNode = findRelatedDeformer(skinObj)
    if not skinNode:
        return
    
    filePath = os.path.join(path, fileName)
    with open(filePath, 'r') as fj:
        weight_dict = json.load(fj)
    skinObj = pm.PyNode(skinNode)
    geo = mc.skinCluster(skinNode, q=True, g=True)[0]
    infJoint = mc.skinCluster(skinNode, q=True, inf=True)
    title = "Be Doing"
    status = 'Sleeping: 0%'
    maxValue = len(infJoint)
    step = 1
    
    star = time.clock()
    
    amount = 0
    mc.progressWindow(title=title,
                      progress=amount,
                      status=status,
                      isInterruptable=True,
                      max=maxValue)
    avgTime = []
    for i, jot in enumerate(infJoint):
        if mc.progressWindow(query=True, isCancelled=True):
            break
        
        # Check if end condition has been reached
        if mc.progressWindow(query=True, progress=True) >= maxValue:
            break
        
        amount += step
        mc.progressWindow(edit=True, progress=amount, status=('Sleeping: ' + `amount` + '/%s' % maxValue))
        
        _star = time.clock()
        
        weight_value = weight_dict['weight']
        wei = weight_value[jot]
        skinObj.setWeights(geo, [i], wei)
        _end = time.clock()
        avgTime.append(_end - _star)
    
    mc.progressWindow(endProgress=1)
    
    # wc.Core().progress(command=command(),
    #                    maxValue=len(joint),
    #                    )
    end = time.clock()
    pm.displayInfo("Weight : Import skinCluster weight file path - {}".format(path))
    pm.displayInfo("Weight : Import skinCluster weight file name - {}".format(fileName + ".json"))
    pm.displayInfo("Weight : Import skinCluster elapsed time - {} s".format(end - star))
    pm.displayInfo("Weight : Import skinCluster average time - {} s".format(sum(avgTime) / len(avgTime)))
    # pm.displayInfo(end-star)
    # print weight_dict
    return


def resetSkinPose(skinObj):
    skinNode = findRelatedDeformer(skinObj)
    if not skinNode:
        return
    skinNode_matrix = "%s.matrix" % skinNode
    matrixList = mc.getAttr(skinNode_matrix, mi=True)
    infs = mc.skinCluster(skinNode, q=True, inf=True)
    if infs:
        for idx in matrixList:
            inf = mc.listConnections("%s[%d]" % (skinNode_matrix, idx), s=True, d=False, scn=True)
            if inf is None:
                continue
            matrix = mc.getAttr("%s.worldInverseMatrix[0]" % inf[0])
            mc.setAttr("%s.pm[%d]" % (skinNode, idx), matrix, type="matrix")
    
    return


def lockInf(skinObj, lock=True, exclude=()):
    skinNode = findRelatedDeformer(skinObj)
    if not skinNode:
        return False
    infJoint = mc.skinCluster(skinNode, q=True, inf=True)
    for jot in infJoint:
        if jot not in exclude:
            mc.setAttr(jot + ".liw", lock)
    for ejot in exclude:
        mc.setAttr(ejot + ".liw", abs(lock - 1))
    return True


def pruneSmallWeight(skinObj, ignore=0.00001):
    if lockInf(skinObj, lock=False):
        skinNode = findRelatedDeformer(skinObj)
        if skinNode:
            mc.select(skinObj, r=True)
            mc.skinPercent(skinNode, prw=ignore)
            return True
        else:
            return False
    else:
        return False


def blendShapeItemExist(blendShape, item):
    tgts = getBlendShapeItemDict(blendShape)
    if item in tgts:
        return True
    else:
        return False


def getBlendShapeIndexDict(blendShape):
    blsDict = dict()
    alias = mc.aliasAttr(blendShape, q=True)
    if alias:
        num = len(alias)
        for i in range(num)[::2]:
            item = alias[i]
            tmp = alias[i + 1]
            index = tmp.split("[")[-1].split("]")[0]
            blsDict[int(index)] = item
    return blsDict


def getBlendShapeItemDict(blendShape):
    blsDict = dict()
    alias = mc.aliasAttr(blendShape, q=True)
    if alias:
        num = len(alias)
        for i in range(num)[::2]:
            item = alias[i]
            tmp = alias[i + 1]
            index = tmp.split("[")[-1].split("]")[0]
            blsDict[item] = int(index)
    return blsDict


def blendShapeRemoveItem(blendShape, item):
    if blendShapeItemExist(blendShape, item):
        index = getBlendShapeItemDict(blendShape)[item]
        baseGeo = mc.blendShape(blendShape, q=True, g=True)[0]
        mc.blendShape(blendShape, e=True, tc=False, rm=True, t=(baseGeo, index, item, 1))
        return True
    else:
        return False


def blendShapeAppendTgt(blendShape, newTgt):
    if not blendShapeItemExist(blendShape, newTgt):
        blsDict = getBlendShapeIndexDict(blendShape)
        if blsDict:
            maxValue = max(blsDict.keys())
        else:
            maxValue = -1
        baseGeo = mc.blendShape(blendShape, q=True, g=True)[0]
        mc.blendShape(blendShape, e=True, t=(baseGeo, maxValue + 1, newTgt, 1))
        return True
    else:
        return False


def blendShapeRenameTgt(blendShape, oldItem, newName):
    if blendShapeItemExist(blendShape, oldItem):
        blsDict = getBlendShapeItemDict(blendShape)
        mc.aliasAttr(newName, "{}.w[{}]".format(blendShape, blsDict[oldItem]))
        return True
    else:
        return False


def getBlendShapeIndexForItem(blendShape, item):
    alias = mc.aliasAttr(blendShape, q=True)
    
    if blendShapeItemExist(blendShape, item):
        itemIndex = alias.index(item)
        tmp = alias[itemIndex + 1]
        index = tmp.split("[")[-1].split("]")[0]
        
        return index
    else:
        return


def getBlendShapeItemTgt(blendShape, item):
    if blendShapeItemExist(blendShape, item):
        idx = getBlendShapeIndexForItem(blendShape, item)
        print idx
        sc = mc.listConnections("{blsNode}.inputTarget[0].inputTargetGroup[{index}]."
                                "inputTargetItem[6000].inputGeomTarget".format(blsNode=blendShape, index=idx),
                                s=True)
        if sc:
            return sc[0]
        else:
            return 
    else:
        return 


def blendShapeReplaceTgt(blendShape, item, newTgt):
    if blendShapeItemExist(blendShape, item):
        newShape = cc.findRelatedShape(newTgt)
        idx = getBlendShapeItemDict(blendShape)[item]
        itemTgt = getBlendShapeItemTgt(blendShape, item)
        tgtAttr = "{blsNode}.inputTarget[0].inputTargetGroup[{index}]." \
                  "inputTargetItem[6000].inputGeomTarget".format(blsNode=blendShape, index=idx)
        
        if itemTgt is not None:
            mc.disconnectAttr("{}.worldMesh[0]".format(itemTgt), tgtAttr)
        
        mc.connectAttr("{}.worldMesh[0]".format(newShape), tgtAttr)
        return True
    else:
        return False


def shapeMirror(origObj, sObj, tObj=None):
    origShape = cc.findRelatedShape(origObj)
    
    tmp = mc.polySphere(name='Temp')[0]
    mc.delete(tmp, ch=True)
    tmpShape = cc.findRelatedShape(tmp)

    mc.connectAttr("{}.worldMesh[0]".format(origShape),
                   "{}.inMesh".format(tmpShape))
    
    if tObj is None:
        tObj = mc.duplicate(tmp, name="{}_mirrorTgt".format(sObj))[0]
    sxobj = mc.duplicate(tmp, name="{}_sx".format(sObj))[0]
    mc.setAttr(sxobj + ".sx", -1)
    mc.select([tmp, sxobj])
    mel.eval("CreateWrap;")
    mc.blendShape(sObj, sxobj, w=[0, 1])
    mc.blendShape(tmp, tObj, w=[0, 1])
    mc.delete([tmp, sxobj])
    
    # mc.disconnectAttr("{}.worldMesh[0]".format(origShape),
    #                   "{}.inMesh".format(tmpShape))
    return tObj


def skinWeightToDeformWeight(sModel, sjoint, tDeform):
    wi = pm.PyNode(tDeform)
    count = mc.polyEvaluate(sModel, v=1)
    
    skinNode = findRelatedDeformer(sModel)
    
    for i in range(count):
        v = mc.skinPercent(skinNode, sModel + '.vtx[%s]' % i, transform=sjoint, query=True)
        if v == 0:
            continue
        wi.wl[0].w[i].set(v)
    
    return


def MoveWireWeightToSkinWeight(crv, wireMode, skinModel, sjoint, moveDis=1, moveAxial="y"):
    crvSkin = findRelatedDeformer(crv)
    crvjoint = mc.skinCluster(crvSkin, q=True, inf=True)
    
    lockInf(skinModel, exclude=[sjoint])
    
    targetSkin = findRelatedDeformer(skinModel)
    targetJoint = mc.skinCluster(targetSkin, q=True, inf=True)
    if not set(crvjoint).issubset(targetJoint):
        njoint = list(set(crvjoint).difference(set(targetJoint)))
        mc.skinCluster(targetSkin, e=True, ai=njoint, lw=True)
    
    orgVtxNumber = mc.polyEvaluate(skinModel, v=True)
    orgVtxtranslate = [mc.xform(skinModel + ".vtx[%s]" % vtx, q=True, ws=True, t=True) for vtx in range(orgVtxNumber)]
    
    _targetJoint = mc.skinCluster(targetSkin, q=True, inf=True)
    
    for cj in crvjoint:
        mc.setAttr(cj + ".t" + moveAxial, moveDis)
        targetVtxtranslate = [mc.xform(wireMode + ".vtx[%s]" % vtx, q=True, ws=True, t=True) for vtx in
                              range(orgVtxNumber)]
        
        for vtx in range(orgVtxNumber):
            if moveAxial == "y":
                # print targetVtxtranslate
                wv = round(abs(targetVtxtranslate[vtx][1] - orgVtxtranslate[vtx][1]), 6)
            
            if wv == 0:
                continue
            orgWvList = mc.skinPercent(targetSkin, skinModel + ".vtx[%s]" % vtx, q=True, v=True)
            
            orgWv = mc.skinPercent(targetSkin, skinModel + ".vtx[%s]" % vtx, q=True, transform=sjoint)
            
            wvlist = []
            for _cj, _wv in zip(_targetJoint, orgWvList):
                if _cj == cj:
                    wvlist.append((_cj, wv))
                elif _cj == sjoint:
                    wvlist.append((_cj, orgWv - wv))
                else:
                    wvlist.append((_cj, _wv))
                # print _cj, _wv
            # print wvlist
            mc.skinPercent(targetSkin, skinModel + ".vtx[%s]" % vtx, transformValue=wvlist)
        
        # return
        mc.setAttr(cj + ".t" + moveAxial, 0)
        print cj
    return


if __name__ == '__main__':
    pass